# encoding: utf-8

__version__ = "1.0"

"""
##--------------------------------------#
## Kvasir
##
## (c) 2010-2014 Cisco Systems, Inc.
## (c) 2015 Kurt Grutzmacher
##
## Exploits utility module
##
## Author: Kurt Grutzmacher <grutz@jingojango.net>
##--------------------------------------#
"""

from gluon import current

import logging
logger = logging.getLogger("web2py.app.kvasir")

##-------------------------------------------------------------------------

def connect_exploits():
    """Connects Vulnerability IDs and Exploits together in t_exploit_references"""

    db = current.globalenv['db']
    cache = current.globalenv['cache']
    exploits_db = db.t_exploits
    exprefs_db = db.t_exploit_references
    vd = db.t_vulndata
    vr = db.t_vuln_refs
    vrefs = db.t_vuln_references
    s_vulndata = db(vr.id == vrefs.f_vuln_ref_id)

    db.t_exploit_references.truncate()
    db.commit()
    existing_vulnids = {}
    for r in db(vd).select(vd.id, vd.f_vulnid, cache=(cache.ram, 60)):
        existing_vulnids[r.f_vulnid] = r.id

    # Run through exploits with vulnid strings set
    exploits = db(exploits_db.f_vulnid != None).select(cache=(cache.ram, 60))
    refs = 0
    for exploit in exploits:
        for vulnid in exploit['f_vulnid']:
            if existing_vulnids.has_key(vulnid):
                refdata = { 'f_vulndata_id': existing_vulnids[vulnid], 'f_exploit_id': exploit.id }
                try:
                    exprefs_db.insert(**refdata)
                    refs += 1
                except:
                    pass
                db.commit()

    existing_cveids = {}
    for r in s_vulndata(vr.f_text != None).select(vr.f_text, vrefs.f_vulndata_id, cache=(cache.ram, 60)):
        existing_cveids[r.t_vuln_refs.f_text] = r.t_vuln_references.f_vulndata_id

    # run through exploits with cve strings set
    for exploit in db(exploits_db.f_cve != None).select(cache=(cache.ram, 60)):
        for cve in exploit['f_cve']:
            if existing_cveids.has_key(cve):
                refdata = { 'f_vulndata_id': existing_cveids[cve], 'f_exploit_id': exploit.id }
                try:
                    exprefs_db.insert(**refdata)
                    refs += 1
                except:
                    pass
                db.commit()

    return refs

##-------------------------------------------------------------------------

def add_exploit(
        cve=None,
        vuln_ids=None,
        f_name=None,
        f_title=None,
        f_description=None,
        f_source=None,
        f_rank='unknown',
        f_level='Unknown'
    ):
    """
    Adds an exploit to the database
    """
    db = current.globalenv['db']

    logging.info(" [*] Adding: %s -- %s" % (f_source, f_title))

    q = (db.t_exploits.f_title == f_title) & (db.t_exploits.f_source == f_source)
    exp_id = db.t_exploits.update_or_insert(
        q,
        f_name=f_name,
        f_title=f_title,
        f_description=f_description,
        f_source=f_source,
        f_rank=f_rank,
        f_level=f_level,
        f_vulnid=vuln_ids,
        f_cve=cve,
    )
    db.commit()

    return exp_id
